from utils import yaml_to_dict,write_file

"""
1.对于历史的问题尽心总结：
    a.单次输入所有的问题，总结效果可能不会好。
    b.迭代输入，每次输入100个问询文件的问题，然后叠加。
2.生成问询问题的列表后，得到具体的问询规则
    a.对历史问询，逐渐输入大模型，得到审核规则
    b.不同批次的审核规则合并
    c.加入文档里面的额外规则
3.加入一些特别有的问询规则
"""

system_prompt = """
你是一名证券交易所的业务专家，在公司上市审核问询领域有何丰富的经验，你的工作职责是对公司上市申请书进行审核，即根据上市申请书的申报材料，提出专业的问询问题'
现在根据你的专业知识，请帮助我完成以下任务。
"""

prompt_template = """
我想利用大模型开发针对公司上市申请文件智能问询的系统，能够对企业在主板、科创板、创业板的上市申请文件进行智能审核，提出对应的专业审核问询问题。
下面是我根据上交所公开的文件，找到的历史上市申请文件问询函里面的论证事项，里面记录了具体的公司名称和问询的问题。
------------------------------【历史问询公司和事项：start】------------------------------
{history}
------------------------------【历史问询公司和事项：end】--------------------------------
现在需要你完成以下工作：
1.将问询问题进行分类，然后列出每个类型下，交易所可能要求发行人、保荐机构和发行人律师论证的事项。
2.针对企业提交新的上市申请书，我们根据此提出智能问询。所以问询问题分类要完整，每种问询分类下面的论证事项也要列完整。
注意，仅仅列出大的问题分类，以及每种分类下面的论证事项，不需要列出论证事项下面的具体问题点。比如：
```
类型：公司基本情况
问题1.历史沿革
问题2：实控人和实控权
问题3：股东和股权结构
```
"""

def llm_etract_prolem_title(problems):
    question_content = prompt_template.format(history=str(problems))
    response = call_with_messages(system_prompt, question_content)
    return response

def batch_extract_problem_title(batch):
    problems = {}
    for company_name, problem_titles in batch:
        problems[company_name] = problem_titles
    res = llm_etract_prolem_title(problems)
    return res


def extract_problem_title(batch_size = 100):
    problem_title_dict = yaml_to_dict('../data/problem_title_filter.txt')
    batch = []
    for i, (key ,value) in enumerate(problem_title_dict.items()):
        if i + 1 <= 2300:
            continue
        batch.append((key, value))
        if (i + 1) % batch_size == 0:
            res = batch_extract_problem_title(batch)
            print(res)
            write_file('problem_categroy/%s.txt' % str(i+1), res)
            batch = []

    if batch:
        res = batch_extract_problem_title(batch)
        write_file('problem_categroy/%s.txt' % str(i + 1), res)



if __name__ == '__main__':
    extract_problem_title()



